/*
 * Copyright 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "pythonutil.h"

#include <iostream>

namespace pythonutil {

std::map<std::string, std::vector<int64_t>>
convertPythonDictToStdMap(PyObject* dictObject) {
    std::map<std::string, std::vector<int64_t>> resultMap;

    if (nullptr == dictObject) {
        return resultMap;
    }

    PyObject *keys = PyDict_Keys(dictObject);
    Py_ssize_t size = PyList_Size(keys);

    for (Py_ssize_t i = 0; i < size; ++i) {
        PyObject *key = PyList_GetItem(keys, i);
        PyObject *value = PyDict_GetItem(dictObject, key);

        if (PyList_Check(value)) {
            std::vector<int64_t> intList;
            Py_ssize_t listSize = PyList_Size(value);
            for (Py_ssize_t j = 0; j < listSize; ++j) {
                PyObject *item = PyList_GetItem(value, j);
                int intValue = PyLong_AsLong(item);
                intList.push_back(intValue);
            }
            resultMap[PyUnicode_AsUTF8(key)] = intList;
        }
    }

    Py_XDECREF(keys);

    return resultMap;
}

std::vector<std::string> convertPythonStringListToStd(PyObject* listObject) {
    std::vector<std::string> resultVector;

    if (nullptr == listObject) {
        return resultVector;
    }

    Py_ssize_t size = PyList_Size(listObject);

    for (Py_ssize_t i = 0; i < size; ++i) {
        PyObject *item = PyList_GetItem(listObject, i);
        if (item == Py_None) {
            continue;
        }
        if (PyUnicode_Check(item)) {
            const char *str = PyUnicode_AsUTF8(item);
            if (str[0] != '\0')
                resultVector.emplace_back(str);
        } else {
            std::cerr << "This Python list does not contain a string type, cancel convert" << std::endl;
            continue;
        }
    }

    return resultVector;
}

std::string convertPythonStringToStd(PyObject* strObject) {
    if (nullptr == strObject) {
        return std::string();
    }
    if (!PyUnicode_Check(strObject)) {
        std::cerr << "Type check error, object isn't string type" << std::endl;
        return std::string();
    }
    const char* utf8Str = PyUnicode_AsUTF8(strObject);
    return std::string(utf8Str);
}

int convertPythonIntToStd(PyObject* intObject) {
    if (nullptr == intObject) {
        return -1;
    }
    if (!PyLong_Check(intObject)) {
        std::cerr << "Type check error, object isn't int type" << std::endl;
        return -1;
    }

    long longValue = PyLong_AsLong(intObject);
    return static_cast<int>(longValue);
}

PyObject* convertStringVectorToPyObject(const std::vector<std::string>& vec)
{
    if (vec.empty()) {
        return Py_None;
    }
    PyObject* list = PyList_New(vec.size());
    if(!list) {
        std::cout << "convertStringVectorToPyObject error!!!" << std::endl;
    }

    for (size_t i = 0; i < vec.size(); ++i) {
        PyObject* pyStr = PyUnicode_DecodeFSDefault(vec[i].c_str());
        if (nullptr == pyStr) {
            Py_XDECREF(list);
            return Py_None;
        }

        PyList_SET_ITEM(list, i, pyStr);
    }
    return list;
}

PyObject* convertDoubleVectorToPyObject(const std::vector<double>& vec) {
    if (vec.empty()) {
         std::cout << "vec is empty" << std::endl;
        return Py_None;
    }
    PyObject* list = PyList_New(vec.size());
    if(!list) {
        std::cout << "convertDoubleVectorToPyObject error!!!" << std::endl;
    }

    for (size_t i = 0; i < vec.size(); ++i) {
        PyList_SET_ITEM(list, i, PyFloat_FromDouble(vec[i]));
    }
    return list;
}

PyObject* convertDoubleMatrixToPyObject(const std::vector<std::vector<double>>& matrix)
{
    if (matrix.empty()) {
        return Py_None;
    }
    PyObject* list = PyList_New(matrix.size());
    if(nullptr == list) {
        std::cout << "convertDoubleMatrixToPyObject error!!!" << std::endl;
    }

    for (size_t i = 0; i < matrix.size(); ++i) {
        if (matrix[i].empty()) {
            return Py_None;
        }
        PyList_SET_ITEM(list, i, convertDoubleVectorToPyObject(matrix[i]));
    }

    return list;
}

// 将std::variant<std::string, int64_t>转换为Python对象
PyObject* convertVariantToPyObject(const std::variant<std::string, int64_t>& var)
{
    if (std::holds_alternative<std::string>(var)) {
        return PyUnicode_FromString(std::get<std::string>(var).c_str());
    } else if (std::holds_alternative<int64_t>(var)) {
        return PyLong_FromLongLong(std::get<int64_t>(var));
    }
    // 如果variant中不是预期的类型，返回None或抛出异常
    return Py_None;
}

// 将std::map<std::string, std::variant<std::string, int64_t>>转换为Python字典对象
PyObject* convertMapVariantToPyObject(const std::map<std::string, std::variant<std::string, int64_t>>& map) {
    PyObject* dict = PyDict_New();
    if (!dict) {
        std::cerr << "Failed to create Python dictionary" << std::endl;
        return Py_None;
    }

    for (const auto& entry : map) {
        PyObject* key = PyUnicode_FromString(entry.first.c_str());
        PyObject* value = convertVariantToPyObject(entry.second);
        if (!key || !value) {
            Py_XDECREF(dict);
            return Py_None;
        }
        PyDict_SetItem(dict, key, value);
        Py_XDECREF(key);
        Py_XDECREF(value);
    }
    return dict;
}

// 将std::vector<std::map<std::string, std::variant<std::string, int64_t>>>转换为Python列表对象
PyObject* convertMapListToPyObject(const std::vector<std::map<std::string, std::variant<std::string, int64_t>>>& mapList) {
    if (mapList.empty()) {
        return Py_None;
    }
    PyObject* list = PyList_New(mapList.size());
    for (size_t i = 0; i < mapList.size(); ++i) {
        PyList_SET_ITEM(list, i, convertMapVariantToPyObject(mapList[i]));
    }
    return list;
}

std::vector<std::vector<double>> convertPythonListToMatrix(PyObject* pyMatrix) {
    std::vector<std::vector<double>> matrix;
    if (nullptr == pyMatrix) {
        return matrix;
    }
    if (!PyList_Check(pyMatrix)) {
        std::cerr << "The matrix object is not a list." << std::endl;
        return matrix;
    }

    if (nullptr == pyMatrix) {
        return matrix;
    }
    Py_ssize_t outer_size = PyList_Size(pyMatrix);
    for (Py_ssize_t i = 0; i < outer_size; ++i) {
        PyObject* inner_list = PyList_GetItem(pyMatrix, i);
        if (PyList_Check(inner_list)) {
            std::vector<double> vector = extractDoubles(inner_list);
            matrix.push_back(vector);
        } else {
            std::cerr << "Non-list value found in the matrix." << std::endl;
        }
    }
    return matrix;
}

std::vector<double> extractDoubles(PyObject* inner_list) {
    std::vector<double> inner_vec;
    if (nullptr == inner_list) {
        return inner_vec;
    }
    Py_ssize_t inner_size = PyList_Size(inner_list);
    for (Py_ssize_t j = 0; j < inner_size; ++j) {
        PyObject* item = PyList_GetItem(inner_list, j);
        if (PyFloat_Check(item)) {
            double value = PyFloat_AsDouble(item);
            inner_vec.push_back(value);
        } else {
            std::cerr << "Non-float value found in the embeddings list." << std::endl;
        }
    }
    return inner_vec;
}

std::vector<std::map<std::string, std::variant<std::string, int64_t>>> convertPythonListToMapOfVariant(PyObject* pyList) {
    std::vector<std::map<std::string, std::variant<std::string, int64_t>>> maoList;
    if (nullptr == pyList) {
        return maoList;
    }
    if (!PyList_Check(pyList)) {
        return maoList;
    }
    Py_ssize_t size = PyList_Size(pyList);
    for (Py_ssize_t i = 0; i < size; ++i) {
        PyObject* item = PyList_GetItem(pyList, i);
        if (PyDict_Check(item)) {
            std::map<std::string, std::variant<std::string, int64_t>> metadata_map = convertPyDictToStdMap(item);
            maoList.push_back(metadata_map);
        }
    }

    return maoList;
}

std::map<std::string, std::variant<std::string, int64_t>> convertPyDictToStdMap(PyObject* dict) {
    std::map<std::string, std::variant<std::string, int64_t>> map;
    if (nullptr == dict) {
        return map;
    }

    PyObject* keys = PyDict_Keys(dict);
    Py_ssize_t keyCount = PyList_Size(keys);
    for (Py_ssize_t j = 0; j < keyCount; ++j) {
        PyObject* pyKey = PyList_GetItem(keys, j);
        if (PyUnicode_Check(pyKey)) {
            std::string key = convertPyUnicodeToStdString(pyKey);
            PyObject* pyValue = PyDict_GetItem(dict, pyKey);
            std::variant<std::string, int64_t> value = convertPyVarToStd(pyValue);
            map[key] = value;
        }
    }
    Py_XDECREF(keys);
    return map;
}

std::string convertPyUnicodeToStdString(PyObject* unicode) {
    if (nullptr == unicode) {
        return std::string();
    }

    PyObject* utf8 = PyUnicode_AsUTF8String(unicode);
    std::string str(PyBytes_AsString(utf8));
    Py_XDECREF(utf8);
    return str;
}

std::variant<std::string, int64_t> convertPyVarToStd(PyObject* pyValue) {
    std::variant<std::string, int64_t> value;
    if (nullptr == pyValue) {
        return value;
    }
    if (PyUnicode_Check(pyValue)) {
        value = convertPyUnicodeToStdString(pyValue);
    } else if (PyLong_Check(pyValue)) {
        int64_t intValue = PyLong_AsLongLong(pyValue);
        value = intValue;
    }
    return value;
}

std::vector<float> convertPythonListToVectorFloat(PyObject* obj)
{
    std::vector<float> result;
    if (nullptr == obj) {
        return result;
    }
    if (!PySequence_Check(obj)) {
        return result;
    }
    Py_ssize_t length = PySequence_Size(obj);
    result.reserve(length);
    for (Py_ssize_t i = 0; i < length; ++i) {
        PyObject* item = PySequence_GetItem(obj, i);
        if (nullptr == item) {
            PyErr_Print();
            return result;
        }
        if (!PyFloat_Check(item)) {
            Py_XDECREF(item);
            return result;
        }
        double value = PyFloat_AsDouble(item);
        if (value == -1.0 && PyErr_Occurred()) {
            PyErr_Print();
            Py_XDECREF(item);
            return result;
        }
        result.push_back(static_cast<float>(value));

        Py_XDECREF(item);
    }
    return result;
}

}
